#include "group_norm_v2_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
namespace optiling {
    constexpr uint32_t ALIGN_NUM = 32;  
    constexpr uint32_t CALC_ALIGN_NUM = 256;  
    constexpr uint32_t BUFFER_NUM = 1;  
    uint32_t RoundUp(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }
        return ((a + b - 1) / b) * b;
    }    
    uint32_t RoundDown(uint32_t a, uint32_t b)
    { 
        if (b == 0) {
            return 0;
        }        
        return (a / b) * b;
    } 
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    GroupNormV2TilingData tiling;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    if (!((socVersion == platform_ascendc::SocVersion::ASCEND910B))) {
        return ge::GRAPH_FAILED;
    }
    uint64_t ub_size;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv();
    const gert::StorageShape* x1_shape = context->GetInputShape(0);
    const gert::Shape shape = x1_shape->GetStorageShape();
    uint32_t totalLength = x1_shape->GetOriginShape().GetShapeSize();
    auto N = shape.GetDim(0);
    auto C = shape.GetDim(1);
    const int64_t* pNum_groups = context->GetAttrs()->GetInt(0);
    const char* pData_format = context->GetAttrs()->GetStr(1); 
    const float* pEpsilon = context->GetAttrs()->GetFloat(2);
    const bool* pIs_training = context->GetAttrs()->GetBool(3);
    auto HW = totalLength / N / C;
    auto NG = N * (* pNum_groups);
    auto C1 = C / (* pNum_groups);
    auto C1HW = totalLength  /N /(* pNum_groups);
    auto dataType = context->GetInputDesc(0)->GetDataType();
    uint32_t dtypeSize;
    uint32_t maxTileLength;
    switch (dataType) {
        case ge::DT_FLOAT:
            maxTileLength = 10240;
            dtypeSize = 4;
            context->SetTilingKey(1);
        break;
        case ge::DT_FLOAT16:
            maxTileLength = 4096;
            dtypeSize = 2;
            context->SetTilingKey(2);
        break;
        default:
            return ge::GRAPH_FAILED;
        break;
    }  
    uint32_t coreNum;
    if(C1HW > coreNum){
        coreNum = aivNum;  
    }
    else
    {
        coreNum = C1HW;
    }
    context->SetBlockDim(coreNum);  
    uint32_t C1HWAlign = RoundUp(C1HW, coreNum);
    uint32_t tileLength = C1HWAlign / coreNum;
    uint32_t tileNum = 1;
    if(tileLength > maxTileLength)
    {
        tileNum = tileLength /  maxTileLength;
        if((tileLength %  maxTileLength) != 0)
        {
            tileNum = tileNum + 1;
        }
        tileLength = maxTileLength;
    }   
    if(HW > coreNum){
        coreNum = aivNum;  
    }
    else
    {
        coreNum = HW;
    }    
    uint32_t HWAlign = RoundUp(HW, coreNum);    
    uint32_t baseLength = HWAlign / coreNum;
    if(baseLength > maxTileLength)
    {
        baseLength = maxTileLength;
    }
    uint32_t tileLengthAlign = RoundUp(tileLength, (ALIGN_NUM / dtypeSize));
    tiling.set_blockLength(totalLength);
    tiling.set_tileNum(tileNum);
    tiling.set_tileLength(tileLength);
    tiling.set_tileLengthAlign(tileLengthAlign);
    tiling.set_baseLength(baseLength);
    tiling.set_NG(NG);
    tiling.set_G((* pNum_groups));
    tiling.set_C1(C1);
    tiling.set_C1HW(C1HW);
    tiling.set_f_nC1HW(-1.0 / (float)(C1HW));
    tiling.set_f_C1HW(1.0 / (float)(C1HW));
    tiling.set_Eps(* pEpsilon);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t usrSize = coreNum * 32; 
    uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = usrSize + sysWorkspaceSize; 
    return ge::GRAPH_SUCCESS;
}
}
namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}
namespace ops {
class GroupNormV2 : public OpDef {
public:
    explicit GroupNormV2(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("gamma")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("beta")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("mean")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("rstd")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("num_groups").Int();
        this->Attr("data_format").AttrType(OPTIONAL).String("NCHW");
        this->Attr("eps").AttrType(OPTIONAL).Float(0.0001);
        this->Attr("is_training").AttrType(OPTIONAL).Bool(true);
        this->SetInferShape(ge::InferShape);
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};
OP_ADD(GroupNormV2);
}
